# coding: utf-8

import os, sys
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../')))
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), './')))
import time
import numpy as np
import logging
import sys
# ##此处使用自己的文件位置
# sys.path.append("D:/teamwork/..")
from XRobot.envs import RobotEnv
from XRobot.commons.cv_utils import get_cam_intrinsic
import XRobot.commons.transform as T
logging.basicConfig(level=logging.INFO)
import cv2

class CatchBlock(object):
    def __init__(self):
        # 定义颜色范围，格式为 (HSV_min, HSV_max)
        # 色调（H）- 色相（S）- 饱和度（V） 下面的颜色范围，可以用图片颜色的直方图来计算得到
        self.color_ranges = {
            'red': [(0, 100, 120), (10, 255, 255)],
            'green': [(35, 50, 50), (85, 255, 255)],
            'blue': [(100, 150, 0), (140, 255, 255)],
            'yellow': [(20, 100, 100), (30, 255, 255)]
        }
        
        self.left_x = 230  # 定义相机坐标系中，机器人底盘中心点的 x 坐标，用于区分左右两块区域
        self.right_x = 300
        self.q = T.euler_2_quat(np.array([0, 0, 90 * np.pi / 180]))  # 最后放置物块时，夹爪的角度是固定的（与相机视野的3d坐标系的y轴保持平行-即平面坐标系的x轴）
        
        
    def find_color_blocks(self, image):
        """检测方形物块的位置

        Args:
            image (np.array): 图片的 numpy 数组

        Returns:
            dict: 四种颜色的物块中心坐标，及物块与x轴的夹角
        """
        # 找到图片中red、green、blue、yellow四种颜色的圆形块，并返回圆心坐标
        # 将图像从 BGR 转换为 HSV
        hsv_image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)

        # 存储结果
        block_positions = {}

        # 遍历每种颜色范围
        for color, (lower, upper) in self.color_ranges.items():
            # 创建掩码
            lower = np.array(lower)
            upper = np.array(upper)
            mask = cv2.inRange(hsv_image, lower, upper)

            # 查找轮廓
            contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
            block_positions[color] = []
            # 遍历轮廓，检测矩形形状
            for contour in contours:
                # 计算轮廓的近似多边形
                epsilon = 0.02 * cv2.arcLength(contour, True)
                approx = cv2.approxPolyDP(contour, epsilon, True)

                # 判断是否为矩形（超过 4 个顶点的都纳入考量）
                if len(approx) >= 4:
                    # 计算矩形的中心点
                    M = cv2.moments(contour)
                    # M['m00'] 是面积
                    if M['m00'] != 0:
                        cx = round(M['m10'] / M['m00'], 1)
                        cy = round(M['m01'] / M['m00'], 1)
                        
                        # 获取最小外接矩形
                        rect = cv2.minAreaRect(contour)
                        
                        # 计算矩形的角度
                        angle = rect[2]
                        
                        # 如果角度为负数，调整角度
                        if angle < -45:
                            angle += 90
                            
                        # M['m00'] 是面积，用面积作为阈值，过滤掉较大的方块（用超过阈值的顶点数，和低于阈值的面积数，来确定方块）
                        if M['m00'] <= 150:
                            block_positions[color].append([round(cx, 1), round(cy, 1), angle])

        return block_positions

    def find_color_circles(self, image):
        """_summary_

        Args:
            image (np.array): 图片的 numpy 数组

        Returns:
            dict: 四种颜色的圆盘的中心坐标
        """
        # 将图像从 BGR 转换为 HSV
        hsv_image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
        
        # 存储结果
        circle_positions = {}

        # 遍历每种颜色范围
        for color, (lower, upper) in self.color_ranges.items():
            # 创建掩码
            lower = np.array(lower)
            upper = np.array(upper)
            mask = cv2.inRange(hsv_image, lower, upper)
            
            # 使用形态学操作去除噪声
            kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (5, 5))
            mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)
            mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)

            # 使用 HoughCircles 检测圆（cv专业检测圆形的方法）
            circles = cv2.HoughCircles(mask, cv2.HOUGH_GRADIENT, dp=1.2, minDist=30,
                                    param1=5, param2=3, minRadius=15, maxRadius=16)

            # 存储圆心坐标
            circle_positions[color] = []

            if circles is not None:
                circles = np.round(circles[0, :]).astype("int")
                for (x, y, r) in circles:
                    center = [round(x, 1), round(y, 1)]
                    # 定义圆盘区域的左右边界（x轴 230~300），超过这个区域的不考虑
                    if x < self.left_x or x > self.right_x:
                        continue
                    circle_positions[color].append(center)

        return circle_positions

    def get_pos(self, point):
        """相机-机器人坐标转换

        Args:
            point (list): 相机照片中的坐标，格式为 [x, y, 1, 1]  # z, w 为 1

        Returns:
            np.array: env内的坐标（世界坐标）
        """
        K = get_cam_intrinsic()  # 相机内参
        T = np.array([[0, 1, 0, 0.5],
                    [-1, 0, 0, 0],
                    [0, 0, 1, 0.8],
                    [0, 0, 0, 1]])  # 相机外参
        K_inv = np.linalg.inv(K)  # 相机内参的逆矩阵
        point_env = np.dot(K_inv, point)  # 像素坐标转换为相机坐标
        point_env = np.dot(T, point_env)  # 相机坐标转换为机器人坐标
        # 等价于 point_env = T @ K_inv @ point  # @ 表示矩阵点乘np.dot()
        return [round(point_env[0], 1), round(point_env[1], 1)]  # 取出 x, y, z 坐标
    
    def deliver_block(self, source_position: list, target_position: list, env: RobotEnv):
        """运送物块

        Args:
            source_position (list): 源位置，格式为 [x, y, angle]  # angle 为方块与x轴的夹角
            target_position (list): 目标位置，格式为 [x, y]
            env (RobotEnv): 机器人环境
        """
        source_point = [source_position[0], source_position[1], 1, 1]
        logging.info(f"相机坐标: {source_point}")
        source_point_robot = self.get_pos(source_point)
        logging.info(f"env-世界坐标: {source_point_robot}")
        q = T.euler_2_quat(np.array([0, 0, (90 - source_position[2]) * np.pi / 180]))  # 根据方块角度，调整机械臂夹爪的角度
        action = np.array([source_point_robot[0], source_point_robot[1], 0.1, q[3], q[0], q[1], q[2]])

        # 运动到物块位置
        for t in range(int(200)):
            env.robot.end['arm0'].open()
            env.step(action)

        # 夹爪向下运动
        action[2] = 0.02
        for t in range(int(100)):
            env.step(action)

        # 夹爪闭合
        for t in range(int(20)):
            env.robot.end['arm0'].close()
            env.step(action)

        # 夹爪向上运动
        action[2] = 0.1
        for t in range(int(100)):
            env.step(action)

        # 运动到对应的圆盘位置
        target_point = [target_position[0], target_position[1], 1, 1]
        target_point_robot = self.get_pos(target_point)
        # 目标位置的夹爪角度是固定的
        action = np.array([target_point_robot[0], target_point_robot[1], 0.1, self.q[3], self.q[0], self.q[1], self.q[2]])
        for t in range(int(200)):
            env.step(action)

        # # 末端旋转
        # action[3:] = np.array([0.707, 0, 0, 0.707])  # [0.707, 0, 0, 0.707] 这个q值的角度，即为末端旋转到90度位置，与self.q的角度一样
        # for t in range(int(100)):
        #     env.step(action)

        # 打开夹爪
        for t in range(int(20)):
            env.robot.end['arm0'].open()
            env.step(action)

        # 回到中间位置
        action[1] = 0
        for t in range(int(100)):
            env.step(action)

def main():
    catch_block = CatchBlock()
    ## TODO: 示例代码，可以删除，并在此编写你的代码
    robot = 'UR5e'
    # robot = 'Panda'
    env = RobotEnv(
        robot = robot + 'Grasp',
        render_mode = 'human',
        control_freq = 200,
        controller = 'CARTIK',
        is_show_camera_in_cv = True,
        is_render_camera_offscreen = True,
        camera_in_render = "0_cam",
        # camera_in_window = "frontview",     # 固定的渲染视角
    )

    start_time = time.time()
    env.reset()
    
    action_init = np.array([0.3, 0.4, 0.4, 1, 0, 0, 0])
    for t in range(int(20)):
        env.step(action_init)

    cv_image = env.render("rgb_array")
    cv2.imwrite(f"image.png", cv_image)
    
    # 各种颜色的源物块的中心坐标和夹角
    block_positions = catch_block.find_color_blocks(cv_image)
    logging.info(f"源物块中心坐标和角度: {block_positions}")
    
    # 各种颜色的目标圆盘的中心坐标
    circle_positions = catch_block.find_color_circles(cv_image)
    logging.info(f"目标圆盘中心坐标: {circle_positions}")
    
    catched_blocks = []  # 已经捕获的块
    planted_positions = []  # 已经种植的位置
    
    for k, v in block_positions.items():
        # k：颜色，v：该颜色的圆形块坐标 （外层循环 k 值，表示物块与圆盘的颜色必须相同，某颜色的物块，只能被放置到相同颜色的圆盘上）
        # 若改变物块颜色，需要相应改变圆盘颜色，比如橙色物块变为红色后，橙色圆盘也需要变成红色（即同时存在两个红色物块，同时存在两个红色圆盘）
        # 若只改变物块颜色，不改变圆盘颜色，需要修改下述的逻辑（且一个圆盘上可能会堆叠2个物块，即：只标记捕获，不标记种植）
        for i in range(len(v)):
            # 只处理位于相机左侧视野的块（避免误识别右边视野中的物体）
            if v[i][0] < catch_block.left_x:
                source_position = v[i]
                if k in circle_positions:
                    for j in range(len(circle_positions[k])):
                        if circle_positions[k][j][0] < catch_block.left_x:
                            # 目标地点在相机左边视野，弃用（只认为相机右侧视野的圆形才是有效的目标）
                            if f"{k}_{j}" not in catched_blocks:
                                planted_positions.append(f"{k}_{j}")
                            continue
                        if circle_positions[k][j][0] > catch_block.left_x:
                            target_position = circle_positions[k][j]
                            # 如果物块未被捕获过，且目标圆盘未被种植过，则捕获并种植
                            if f"{k}_{i}" not in catched_blocks and f"{k}_{j}" not in planted_positions:
                                catch_block.deliver_block(source_position, target_position, env)
                                catched_blocks.append(f"{k}_{i}")  # 标记为捕获
                                planted_positions.append(f"{k}_{j}")  # 标记为种植

    task_time = time.time() - start_time
    logging.info("time: %.2f" % task_time)

    env.close()


if __name__ == "__main__":
    main()



